/*
 * Translation to C from nbody0 example in Sverre Aarseth's website.
 */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#include <cuda.h>

#include "ffdot.h"


#define min(x,y) (x<y ? x : y)

float *X, *XDOT, *F, *FDOT, *auxF, *auxFDOT, *D2, *D3;
float4 *X_d, *XDOT_d;
float3 *F_d, *FDOT_d;

float *X0,*X0DOT,*T0,*STEP,firr[3],fd[3];

int    *NEXT;
int    *NEXT_d;

float eta,dtnext,tcrit,eps2;
long Np;
float Time=0.0,tnext=0.0,smax=64.0,smin=pow(2,-16), tmin=1e10,nsteps=0;

void allocate_arrays( long N )
{
        X     = (float *)malloc( sizeof(float)*4 * N );
        X0    = (float *)malloc( sizeof(float)*4 * N );
        XDOT  = (float *)malloc( sizeof(float)*4 * N );
        X0DOT = (float *)malloc( sizeof(float)*4 * N );
        F     = (float *)malloc( sizeof(float)*3 * N );
        FDOT  = (float *)malloc( sizeof(float)*3 * N );
        
        auxF    = (float *)malloc( sizeof(float)*3 * N );
        auxFDOT = (float *)malloc( sizeof(float)*3 * N );
        
        
        //M     = (float *)malloc( sizeof(float)*3 * N );
        D2    = (float *)malloc( sizeof(float)*3 * N );
        D3    = (float *)malloc( sizeof(float)*3 * N );
        
        T0   = (float *)malloc( sizeof(float) * N );
        STEP = (float *)malloc( sizeof(float) * N );
        
        NEXT = (int    *)malloc( sizeof(int)    * N );
        
        cudaMalloc( (void **) &X_d    , sizeof(float4)*N );
        cudaMalloc( (void **) &XDOT_d , sizeof(float4)*N );
        cudaMalloc( (void **) &F_d    , sizeof(float3)*N );
        cudaMalloc( (void **) &FDOT_d , sizeof(float3)*N );
        
        cudaMalloc( (void **) &NEXT_d , sizeof(int   )*N );
        

}

void read_ic( )
{
        FILE *f;
        f = fopen( "../IC/init.txt", "r" );
    
        fscanf( f, "%ld%f%f%f%f\n", &Np, &eta, &dtnext, &tcrit, &eps2 );
        
        allocate_arrays( Np );
        
        for( int j=0; j<Np; j++ )
        {       
                fscanf( f, "%f", &X[4*j+3] );
                
                for( int k=0; k<3; k++ )
                        fscanf( f, "%f", &X[4*j+k] );
        
                for( int k=0; k<2; k++ )
                        fscanf( f, "%f", &XDOT[4*j+k] );
        
                fscanf( f, "%f\n", &XDOT[4*j+2]);
        
        }
        
        fclose( f );

}

void dump_phasecoords( FILE *phstream )
{
        for( long j=0; j<Np; j++ )
        {
                fprintf( phstream, "%f ", X[4*j+3] );
                
                for( int k=0; k<3; k++ )
                        fprintf( phstream, "%f ",  X[4*j+k] );
                
                for( int k=0; k<2; k++ )
                        fprintf( phstream, "%f ", XDOT[4*j+k] );
                
                fprintf( phstream, "%f\n", XDOT[4*j+2]);
        }
}

void FFDOT( long i, float *firr, float *fd, long N )
{

        float DR[3];
        float DV[3];
        
        float RIJ2,DR3I,DRDV;
        
        for( int k=0; k<3; k++ )
        {
                firr[k] = 0.0;
                fd[k] = 0.0;
        }
        
        for( long j=0; j<N; j++ )
        if( j != i )
        {
                for( int k=0; k<3; k++ )
                {
                        DR[k] =    X[4*j+k] -    X[4*i+k];
                        DV[k] = XDOT[4*j+k] - XDOT[4*i+k];
                }
                
                RIJ2 = DR[0]*DR[0] + DR[1]*DR[1] + DR[2]*DR[2] + eps2;
                DR3I = X[4*j+3]/( RIJ2*sqrt(RIJ2) );
                
                DRDV = 3.0*( DR[0]*DV[0] + DR[1]*DV[1] + DR[2]*DV[2] )/RIJ2;
                
                for( int k=0; k<3; k++ )
                {
                        firr[k] +=   DR[k]*DR3I;
                        fd[k] += ( DV[k] - DR[k]*DRDV )*DR3I;
                }  
        }

}

void dump_energy( FILE *Estream )
{
        float K = 0.0,U = 0.0;
        
        for( long i=0; i<Np; i++ )
        {
                K += X[4*i+3]*(pow(XDOT[4*i+0],2)+pow(XDOT[4*i+1],2)+pow(XDOT[4*i+2],2));
                
                for( long j=i+1; j<Np; j++ )
                        U += X[4*i+3]*X[4*j+3]/sqrt( pow(X[4*j+0]-X[4*i+0],2) +
                                                                         pow(X[4*j+1]-X[4*i+1],2) + 
                                                                         pow(X[4*j+2]-X[4*i+2],2) + eps2 ); 
        }
        
        fprintf( Estream, "%f %f %3.7e\n", Time, (float)nsteps, 0.5*K-U);
}


int main( void )
{
	FILE *ph;
	FILE  *E;
	
	ph = fopen( "../output/coord.txt", "w" );
	 E = fopen( "../output/energ.txt", "w" );

	//Read initial conditions
	read_ic();
		
    cuda_ffdot( X, XDOT, auxF, auxFDOT, NEXT, eps2, Np, X_d, XDOT_d, F_d, FDOT_d, NEXT_d );
    
	//Initialize time steps
	for( int j=0; j<Np; j++ )
	{
		//FFDOT(j, firr, fd, Np);
		
		for( int k=0; k<3; k++ )
		{
			firr[k] = auxF[3*j+k];
			  fd[k] = auxFDOT[3*j+k];
		}
	
		for( int k=0; k<3; k++ )
		{
			F[3*j+k]    = 0.5*firr[k];
			FDOT[3*j+k] = fd[k]/6.0;
		}
		for( int k=0; k<4; k++ )
		{
			X0[4*j+k] = X[4*j+k];
			X0DOT[4*j+k] = XDOT[4*j+k];
		}
		
		STEP[j] = sqrt(eta/sqrt(firr[0]*firr[0]+firr[1]*firr[1]+firr[2]*firr[2]));
		
		int K = 2 + log(smax/STEP[j])/log(2.0);
		
		STEP[j] = smax/pow(2,K-1);
		tmin = min(STEP[j], tmin);
		
		T0[j] = 0.0;
	}
	
	//Evolve while
	while( Time < tcrit )
	{
	
	if( Time >= tnext )
	{
		dump_phasecoords( ph );
		//dump_energy     (  E );
	
		fprintf( stderr, "%lf  %lf\n", Time, nsteps );
	
		tnext += dtnext;
	}
	
    // Determines the next block of particles to be advanced.
	
	int length = 0;
	for( int j=0; j<Np; j++ )
	{
		NEXT[j] = 0;
		
		//Advance this particle
		if( T0[j] + STEP[j] == tmin )
		{   
			NEXT[j] = 1;
            length++;
		}
	}

	Time = tmin;
	
	//Predict all coordinates to order FDOT

	for( int j=0; j<Np; j++ )
	{
		float s = Time - T0[j];
		for( int k=0; k<3; k++)
		{   
			   X[4*j+k] = ( ( FDOT[3*j+k]*s + F[3*j+k] )*s + X0DOT[4*j+k] )*s + X0[4*j+k];
		
			XDOT[4*j+k] = ( FDOT[3*j+k]*1.5*s + F[3*j+k] )*2.0*s + X0DOT[4*j+k];
		}
	}
	
	tmin = 1.0e10;
	
	float alpha = 0.05;
	
	char FLAG = length < alpha * Np ? 'C' : 'G';
	
	//printf( "%e %d %cPU\n", Time, length, FLAG );
	
	if( FLAG == 'G' )
		cuda_ffdot( X, XDOT, auxF, auxFDOT, NEXT, eps2, Np, X_d, XDOT_d, F_d, FDOT_d, NEXT_d );
    
	// Advance the next block of particles
	int i;
	for( i=0; i<Np; i++ )
	if( NEXT[i] == 1)
	{
		if( FLAG == 'G' )
		for( int k=0; k<3; k++ )
		{
			firr[k] = auxF[3*i+k];
			  fd[k] = auxFDOT[3*i+k];
		}
		else
			FFDOT(i, firr, fd, Np );
		
	 // Build the Hermite 4th order corrector
		
		float dt = Time - T0[i];
		T0[i] = Time;
		for( int k=0; k<3; k++ )
		{
			float df  = 2.0*F[3*i+k] - firr[k];
			float fid = 6.0*FDOT[3*i+k];
			float sum = fid + fd[k];
			
			float at3 =  2.0*df + dt*(sum      );
			float bt2 = -3.0*df - dt*(sum + fid);
			
			//Corrector step
			X0[4*i+k] = (0.6*at3 + bt2)*dt*dt/12.0 + X[4*i+k];
			
			X0DOT[4*i+k] = (0.75*at3 + bt2)*dt/3.0 + XDOT[4*i+k];     
			
			F[3*i+k]    = 0.5*firr[k];
			FDOT[3*i+k] = fd[k]/6.0;
			
			//New timestep array
			D2[3*i+k] = (3.0*at3 + bt2)*2.0/(dt*dt);
			D3[3*i+k] = at3*6.0/(dt*dt*dt);  
		}
		
		//New timestep calculation
		
		float fi2  = pow(firr[0],2) + pow(firr[1],2) + pow(firr[2],2);
		float fd2  = pow(fd[0],2) + pow(fd[1],2) + pow(fd[2],2);
		float fd22 = pow(D2[3*i+0],2) + pow(D2[3*i+1],2) + pow(D2[3*i+2],2);
		float fd32 = pow(D3[3*i+0],2) + pow(D3[3*i+1],2) + pow(D3[3*i+2],2);
		
		float si = sqrt(eta*(sqrt(fi2*fd22) + fd2)/(sqrt(fd2*fd32) + fd22));
		
		if( si > 2.0*STEP[i] && fmod( (double)Time, 2.0*STEP[i]) == 0.0)
			STEP[i] = min(2.0*STEP[i],smax);
		if( si < STEP[i] && si > smin )
			STEP[i] *= 0.5;
		
		tmin = min(T0[i]+STEP[i],tmin);
		nsteps += 1;
	}
	
	}
	
	fclose( E);
	fclose(ph);
	
	cudaFree(X_d);
	cudaFree(XDOT_d);
	cudaFree(F_d);
	cudaFree(FDOT_d);
	cudaFree(NEXT_d );
	
	return 0;    
}



